from PIL import Image, ImageDraw
import face_recognition
import numpy
import matplotlib.path as mpath

def same(a,b,thr = 50):

    aa = numpy.array([255,255,255],dtype=numpy.int32)
    bb = numpy.array([255,255,255],dtype=numpy.int32)

    for idx, j in enumerate(a):
        aa[idx] = j

    for idx, j in enumerate(b):
        bb[idx] = j

    norm =  numpy.linalg.norm(aa-bb)

    return ( norm < thr )

def gethead(outline, centerh, chinPts):

    head = []

    for pts in outline[0]:
        if pts[1] > centerh :
            break
        else:
            head.append(pts)

    for pts in chinPts:
        if pts[1] >= centerh :
            head.append(pts)

    for pts in outline[1]:
        if pts[1] < centerh :
            head.append(pts)

    return head

def store(src, head):

    min_h = min_w = 0xffffff
    max_h = max_w = -1

    for pts in head:

        if pts[0] < min_w:
            min_w = pts[0]

        if pts[0] > max_w:
            max_w = pts[0]

        if pts[1] < min_h:
            min_h = pts[1]

        if pts[1] > max_h:
            max_h = pts[1]

    box = (min_w,min_h,max_w+1,max_h+1)

    headimg = src.crop(box)

    maskimg = Image.new('L',headimg.size,0)

    path = mpath.Path(head)

    for h in range( min_h, max_h + 1, 1 ):
        for w in range(min_w, max_w + 1, 1):
            if path.contains_point((w,h)) == False:
                maskimg.putpixel( (w-min_w,h-min_h), 0 )
            else:
                maskimg.putpixel( (w-min_w,h-min_h), 255 )

    return headimg, box, maskimg

bgclr = (255,255,255)

path = "/Users/vista/Documents/PIC/searchimg/model"

# Load the jpg file into a numpy array
image = face_recognition.load_image_file(path + "/4.jpeg")

img2 = face_recognition.load_image_file(path + "/12.jpeg")

face_landmarks_list = face_recognition.face_landmarks(image)

face_landmarks_list2 = face_recognition.face_landmarks(img2)

if len(face_landmarks_list) != 1 or len(face_landmarks_list2) != 1:

    print ("face number is not always be 1")

    exit(0)

# face_landmarks = face_landmarks_list[0]

# face_landmarks2 = face_landmarks_list2[0]


bgImg = image.copy()
# value = np.array([255,255,255],dtype=np.uint8)
bgImg.fill(255)


bgImg2 = img2.copy()

bgImg2.fill(255)

# pil_bg_image = Image.fromarray(bgImg)


def getoutline(img, outline):

    for h in range( 0, img.shape[0] , 10):
        for w in range( 0 , img.shape[1], 1):
            if( same( img[h,w] , bgImg[h,w] ) == False ):
                val = (w,h)
                outline[0].append(val)
                break

    for h in range(img.shape[0]-10, 0, -10):
        for w in range(img.shape[1]-1, 0, -1):
            if (same(img[h, w], bgImg[h, w]) == False):
                val = (w, h)
                outline[1].append(val)
                break

    return outline

outline = getoutline(image,[[],[]])

outline2 = getoutline(img2,[[],[]])

# Find all facial features in all the faces in the image


def drawhead(img, face_landmarks, outline):

    pil_image = Image.fromarray(img)

    # d = ImageDraw.Draw(pil_image, 'RGBA')

    chinpts = face_landmarks['chin']

    centerh = 0

    left = 0xfffff
    right = 0
    bottom = 0

    for pts in chinpts:
        centerh += pts[1]
        if pts[1]>bottom:
            bottom = pts[1]
        if pts[0] < left:
            left = pts[0]
        if pts[0] > right:
            right = pts[0]

    centerh /= len(chinpts)

    head = gethead(outline, centerh, chinpts)

    headimg, box, mask = store(pil_image, head)

    # d.line( head , fill=(0, 250, 0, 80), width=4)
    # d.polygon(outline, fill=(255, 0, 255, 80))

    # pil_image.save(path+"/ss.jpeg")

    # headimg.show()

    return headimg, box, pil_image, mask, (bottom,left,right)

headimg, box, pilimage, mask, (bottom, left, right) = drawhead(image, face_landmarks_list[0], outline)

headimg2, box2, pilimage2, mask2, (bottom2, left2, right2) = drawhead(img2, face_landmarks_list2[0], outline2)

newsize = ( (right2 - left2 + 1) * (box[2] - box[0]) / (right - left + 1),
            (right2 - left2 + 1) * (box[3] - box[1]) / (right - left + 1) )

centerw = (right2+left2)/2

nsImg = headimg.resize(newsize)

mask = mask.resize(newsize)
# mask = mask.resize((box2[2] - box2[0],box2[3] - box2[1]))

# size = (box2[2]-box2[0],box2[3]-box2[1]+10)
#
# eraseimg = Image.new(mode='RGB',size=size,color = (255,255,255))

# eraseimg.show()

# clear background of the replace image
# pilimage2.paste((255,255,255), (box2[0],box2[1]-10,box2[2],box2[3]))

# r,g,b = mask.split()

h0 = box2[3] - newsize[1]
w0 = box2[2] - newsize[0]

pilimage2.paste( nsImg , (w0,h0), mask)
# d.bitmap((box2[0],box2[1]),headimg,fill=(255,0,0,80))
# d.bitmap((box2[0], box2[1]), pil_image, fill=(255,255,255))

pilimage2.show()
